import io

import torch
from torchvision import transforms
from PIL import Image
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from timm import create_model
import logging

logging.basicConfig(filename="test.log",level=logging.INFO,format="%(asctime)s - %(levelname)s - %(message)s")

# 创建 FastAPI 应用实例
app = FastAPI()
#  uvicorn test-02:app --reload 启动项目

# 初始化模型
model = create_model("efficientnet_b0", pretrained=False, num_classes=17)  # num_classes 替换为你的类别数
# 加载训练好的模型权重
model.load_state_dict(torch.load('efficientnet_b0_custom.pth',map_location=torch.device('cpu')))

# 将模型设置为评估模式
model.eval()

# 定义数据预处理步骤，使输入图像符合模型要求
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


# 定义预测接口
@app.post("/predict")
async def predict(image: UploadFile = File(...)):
    logging.info(">>>>>>>>>>>>开始预测图片")
    try:
        print(image)
        # 读取上传的图像文件
        contents = await image.read()
        # 将二进制内容转换为类文件对象
        img = Image.open(io.BytesIO(contents))
        print("开始1")
        # 对图像进行预处理
        input_tensor = preprocess(img)
        # 添加一个维度以匹配模型输入要求
        input_batch = input_tensor.unsqueeze(0)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # 如果有可用的 GPU，将输入数据和模型移动到 GPU 上
        if torch.cuda.is_available():
            input_batch = input_batch.to(device)
            model.to(device)

        # 进行推理
        with torch.no_grad():
            output = model(input_batch)
            probabilities = torch.nn.functional.softmax(output, dim=1)
            print(probabilities)
        _, predicted_idx = torch.max(output, 1)
        class_names = ['丝瓜', '人参果', '佛手瓜', '冬瓜', '南瓜', '哈密瓜', '木瓜', '甜瓜-伊丽莎白', '甜瓜-白',
                       '甜瓜-绿', '甜瓜-金', '白兰瓜', '羊角蜜', '苦瓜', '西瓜', '西葫芦', '黄瓜']  # 替换为你的类别名称
        logging.info(">>>>>>>>>>>>预测结果",predicted_idx)
        logging.info(">>>>>>>>>>>>预测结束")
        # 返回预测结果的 JSON 响应
        # 计算概率值
        probabilities = torch.nn.functional.softmax(output, dim=1)
        probabilities = probabilities.numpy()[0].tolist()
        # 获取预测类别
        predicted_class = torch.argmax(output, dim=1).item()
        # 返回预测结果和概率值
        return {"predicted_class": predicted_class, "probabilities": probabilities,"class_name": class_names[predicted_idx]}
    except Exception as e:
        # 若出现异常，返回错误信息和 500 状态码
        return JSONResponse(content={"error": str(e)}, status_code=500)
